"""
内容：训练过程
日期：2020年6月28日
作者：Howie
"""
# 0. 调用要使用的包
from keras.utils import np_utils
from keras.models import Sequential, load_model
from keras.layers import Dense
from keras.utils import plot_model
from keras.callbacks import TensorBoard, EarlyStopping
import numpy as np
import matplotlib.pyplot as plt
import os
from datetime import datetime
# 预设
os.environ["PATH"] += os.pathsep + 'F:/Anaconda3/Grahpviz/bin'  # 添加环境变量
LOGS_SAVE_PATH = './logs/Acc-Loss Curve(' + str(
    datetime.now().strftime('%Y-%m-%d-%H-%M')) + ').pdf'  # 训练讯息保存路径
GRAPH_SAVE_PATH = './graph/'  # TensorBoard
tb_hist = TensorBoard(
    log_dir=GRAPH_SAVE_PATH,
    histogram_freq=0,
    write_graph=True,
    write_images=True)
early_stopping = EarlyStopping(
    monitor='val_accuracy',
    patience=20
)  # 设置早停，监控每个训练周期的验证精度
CALL_BACK_FUNC = [tb_hist, early_stopping]  # 回调函数
DATASET_PATH = '../dataset/mnist/mnist.npz'  # 数据集路径
MODEL_PATH = './infer_model/'  # 模型保存路径
MODEL_VISUAL = './logs/Model.pdf'  # 模型可视化

# 超参
N_SAMPLES_TRAIN = 60000     # 训练集样本数量
N_SAMPLES_TEST = 10000  # 测试集样本数量
IMG_SIZE = 28  # 图片尺寸
UNITS_1 = 64  # 隐藏层1神经元个数
N_CLS = 10  # 类别数
BATCH_SIZE = 64  # 批次数量
EPOCH = 25  # 周期
VAL_SIZE_RATE = 0.3  # 验证集比例


def hist_plot(hist):
    """
    # 可视化训练过程
    :param hist: history对象
    :return:
    """
    fig, loss_ax = plt.subplots()
    acc_ax = loss_ax.twinx()
    # 每个训练周期的训练误差与验证误差
    loss_ax.plot(hist.history['loss'], 'y', label='train loss')
    loss_ax.plot(hist.history['val_loss'], 'r', label='val loss')
    # 每个训练周期的训练精度与验证精度
    acc_ax.plot(hist.history['accuracy'], 'b', label='train acc')
    acc_ax.plot(hist.history['val_accuracy'], 'g', label='val acc')
    # 横轴与纵轴
    loss_ax.set_xlabel('epoch')
    loss_ax.set_ylabel('loss')
    acc_ax.set_ylabel('accuracy')
    # 标签
    loss_ax.legend(loc='upper left')
    acc_ax.legend(loc='lower left')
    # 保存
    plt.savefig(LOGS_SAVE_PATH)
    # 展示
    plt.show()


def load_dataset():
    """
    # 生成数据集
    :return: 划分好的训练集和测试集
    """
    with np.load(DATASET_PATH) as data:
        (X_train, Y_train), (X_test, Y_test) = (
            data['x_train'], data['y_train']), (data['x_test'], data['y_test'])
    # 数据集预处理
    X_train = X_train.reshape(N_SAMPLES_TRAIN,
                              IMG_SIZE * IMG_SIZE).astype('float32') / 255.0
    X_test = X_test.reshape(N_SAMPLES_TEST,
                            IMG_SIZE * IMG_SIZE).astype('float32') / 255.0
    # 标签数据独热编码处理
    Y_train = np_utils.to_categorical(Y_train)
    Y_test = np_utils.to_categorical(Y_test)
    # 打印划分画好的数据集规模
    print('train samples: {}\n'
          'test samples: {}'.format(X_train.shape[0], X_test.shape[0]))

    return X_train, Y_train, X_test, Y_test


def model_building():
    """
    # 模型搭建
    :return: 搭建好的模型
    """
    model = Sequential()
    model.add(
        Dense(
            units=UNITS_1,
            input_dim=IMG_SIZE *
            IMG_SIZE,
            activation='relu'))
    model.add(Dense(units=N_CLS, activation='softmax'))
    plot_model(model, to_file=MODEL_VISUAL, show_shapes=True)  # 深度学习模型可视化
    return model


def train():
    """
    # 训练与评判模型
    :return:
    """
    """训练部分：训练传感与数据集生成"""
    model = model_building()
    X_train, Y_train, X_test, Y_test = load_dataset()
    """训练部分：深度学习模型训练"""
    # 设置训练过程
    model.compile(
        loss='categorical_crossentropy',
        optimizer='sgd',
        metrics=['accuracy'])
    # 训练模型
    hist = model.fit(
        X_train,
        Y_train,
        epochs=EPOCH,
        batch_size=BATCH_SIZE,
        validation_split=VAL_SIZE_RATE,
        callbacks=CALL_BACK_FUNC)
    # 保存已训练模型
    model.save(os.path.join(MODEL_PATH, 'mnist_mlp_model.h5'))
    """评判过程：评判传感"""
    # 查看训练过程
    hist_plot(hist)
    # 模型评价
    model_evaluation(model, X_test, Y_test)
    """评判过程：深度学习模型评判"""
    # 使用模型
    model_application(
        os.path.join(
            MODEL_PATH,
            'mnist_mlp_model.h5'),
        X_test,
        Y_test)


def model_evaluation(model, X_test, Y_test):
    """
    # 模型评价
    :param model: 模型
    :param X_test: 测试集样本
    :param Y_test: 测试集标签
    :return:
    """
    loss_and_metrics = model.evaluate(X_test, Y_test, batch_size=BATCH_SIZE)
    print('Loss and metrics: {}'.format(str(loss_and_metrics)))


def model_application(infer_model_path, X_test, Y_test):
    """
    # 使用模型
    :param infer_model_path: 模型保存路径
    :param X_test: 测试集样本
    :param Y_test: 测试集标签
    :return:
    """
    # 调用模型
    model = load_model(infer_model_path)
    # 准备实操中要使用的真实数据
    X_hat_idx = np.random.choice(X_test.shape[0], 5)
    X_hat = X_test[X_hat_idx]
    # 使用模型
    Y_hat = model.predict_classes(X_hat)

    for i in range(5):
        print('True: ' +
              str(np.argmax(Y_test[X_hat_idx[i]])) +
              ', Predict: ' +
              str(Y_hat[i]))


if __name__ == '__main__':
    train()
